package com.pig4cloud.pigx.knowledge.service.impl;

import com.pig4cloud.pigx.knowledge.service.EmbeddingStoreService;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.qdrant.QdrantEmbeddingStore;
import io.qdrant.client.QdrantClient;
import io.qdrant.client.grpc.Points;
import lombok.RequiredArgsConstructor;
import lombok.SneakyThrows;
import org.springframework.stereotype.Service;

import java.util.*;

import static io.qdrant.client.PointIdFactory.id;

/**
 * @author lengleng
 * @date 2024/9/26
 * <p>
 * 根据不同的集合名称，获取不同的嵌入存储
 */
@Service
@RequiredArgsConstructor
public class EmbeddingStoreServiceImpl implements EmbeddingStoreService {

	private final static Map<String, EmbeddingStore<TextSegment>> EMBEDDING_STORE_MAP = new HashMap<>();

	private final QdrantClient qdrantClient;

	/**
	 * 嵌入存储
	 * @param collectionName 集合名称
	 * @return {@link EmbeddingStore }<{@link TextSegment }>
	 */
	@Override
	public EmbeddingStore<TextSegment> embeddingStore(String collectionName) {
		EmbeddingStore<TextSegment> embeddingStore = EMBEDDING_STORE_MAP.get(collectionName);
		if (Objects.isNull(embeddingStore)) {
			embeddingStore = new QdrantEmbeddingStore(qdrantClient, collectionName, "doc_content");
			EMBEDDING_STORE_MAP.put(collectionName, embeddingStore);
		}

		return embeddingStore;
	}

	/**
	 * 删除
	 * @param collectionName 集合名称
	 * @param idList ID 列表
	 */
	@Override
	@SneakyThrows
	public void delete(String collectionName, List<String> idList) {
		List<Points.PointId> ids = idList.stream().map(id -> id(UUID.fromString(id))).toList();
		this.qdrantClient.deleteAsync(collectionName, ids).get();
	}

}
